#!/usr/bin/env python3
# set coding: utf-8
"""
@Project : ywcmdb
@File :    celery.py
@Time :    2025/1/12 11:59
@Author :  richard zhu
@Email :   gaotao.zhugt@dtzhejiang.com
@purpose : 
"""

from __future__ import absolute_import, unicode_literals
import os
from celery import Celery
from celery.signals import task_postrun
import functools
import logging

# set the default Django settings module for the 'celery' program.
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'ywadmin.settings')

from celery import platforms
from django.conf import settings

app = Celery('tasks')

app.config_from_object(settings, namespace='CELERY')

app.conf.enable_utc = settings.USE_TZ

app.autodiscover_tasks(lambda: settings.INSTALLED_APPS)

logger = logging.getLogger("tasks")


@app.task(bind=True)
def debug_task(self):
    logger.debug(f'Request: {self.request!r}')


def retry_base_task_error():
    """
    celery 失败重试装饰器
    :return:
    """

    def wraps(func):
        @app.task(bind=True, retry_delay=180, max_retries=3)
        @functools.wraps(func)
        def wrapper(self, *args, **kwargs):
            try:
                return func(*args, **kwargs)
            except Exception as exc:
                raise self.retry(exc=exc)
        return wrapper
    return wraps

@task_postrun.connect
def add_periodic_task_name(sender, task_id, task, args, kwargs, **extras):
    """
    任务执行完成后的处理函数
    
    参数说明：
    - sender: 发送信号的任务
    - task_id: 任务ID
    - task: 任务对象
    - args: 任务的位置参数
    - kwargs: 任务的关键字参数
    - retval: 任务的返回值
    - state: 任务的状态（成功/失败）
    """
    state = extras.get('state')
    logger.debug(f"任务 {task.name}[{task_id}] 执行完成，状态: {state}")
    
    # 从任务请求上下文中获取 periodic_task_name
    periodic_task_name = None
    
    # 1. 尝试从 task.request 的顶层字段获取（PeriodicTask 的 headers 会被展开到顶层）
    if hasattr(task, 'request'):
        # task.request 是一个 Context 对象，可以像字典一样访问
        if hasattr(task.request, 'get'):
            periodic_task_name = task.request.get('periodic_task_name')
        # 如果是属性访问
        elif hasattr(task.request, 'periodic_task_name'):
            periodic_task_name = task.request.periodic_task_name
    
    # 2. 如果没有，尝试从 kwargs 中获取（兼容手动执行的任务）
    if not periodic_task_name:
        periodic_task_name = kwargs.get('periodic_task_name')
    

    if periodic_task_name:
        from django_celery_results.models import TaskResult
        # 更新 TaskResult 表中的 periodic_task_name 字段
        updated_count = TaskResult.objects.filter(task_id=task_id).update(periodic_task_name=periodic_task_name)
        logger.debug(f"已更新 TaskResult periodic_task_name: {periodic_task_name}, 影响行数: {updated_count}")
    else:
        logger.debug(f"警告: 未找到 periodic_task_name，任务ID: {task_id}")
    
    return 

